import torch.optim
from dataloader import SelfDataSet
from log import Logger
from plot import plot_picture
import os
import torch.nn as nn
import sys
from unet import UNet
from torch.utils.data import Dataset
from torch import optim, utils
import time

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
net = UNet(n_channels=3, n_classes=1,bilinear=False)
# 将网络拷贝到deivce中
net.to(device=device)
# 加载模型参数
checkpoint = torch.load('train_ape_linemod',map_location=device)
net.load_state_dict(checkpoint['net'])
# 测试模式
net.eval()
loss_fun = nn.BCEWithLogitsLoss()
# 读取所有图片路径
Test_Data_path = '/data/dataset/cdpn/lm_full/real_test/ape/'
tests_path = glob.glob(Test_Data_path + '*-color.png')

class FeatureExtractor():
    """ Class for extracting activations and
    registering gradients from targetted intermediate layers """

    def __init__(self, model, target_layers):
        self.model = model
        self.target_layers = target_layers
        self.gradients = []

    def save_gradient(self, grad):
        self.gradients.append(grad)

    def __call__(self, x):  # 在对x求导时，将梯度保存下来
        outputs = []
        self.gradients = []
        for name, module in self.model._modules.items():
            x = module(x)
            if name in self.target_layers:
                x.register_hook(self.save_gradient)
                outputs += [x]
        return outputs, x


class ModelOutputs():
    """ Class for making a forward pass, and getting:
    1. The network output.
    2. Activations from intermeddiate targetted layers.
    3. Gradients from intermeddiate targetted layers. """

    def __init__(self, model, feature_module, target_layers):
        self.model = model
        self.feature_module = feature_module
        self.feature_extractor = FeatureExtractor(self.feature_module, target_layers)

    def get_gradients(self):
        return self.feature_extractor.gradients

    def __call__(self, x):
        target_activations = []
        # for name, module in self.model._modules.items():
        #     print(module)
        #     if module == self.feature_module:
        #         target_activations, x = self.feature_extractor(x)
        #     elif "avgpool" in name.lower():
        #         x = module(x)
        #         x = x.view(x.size(0), -1)
        #     else:
        #         print(x.shape)
        #         x = module(x)
        x=model(x)
        #print(x.shape)
        target_activations, x = self.feature_extractor(x)



        return target_activations, x


def preprocess_image(img):
    means = [0.485, 0.456, 0.406]
    stds = [0.229, 0.224, 0.225]

    preprocessed_img = img.copy()[:, :, ::-1]
    for i in range(3):
        preprocessed_img[:, :, i] = preprocessed_img[:, :, i] - means[i]
        preprocessed_img[:, :, i] = preprocessed_img[:, :, i] / stds[i]
    preprocessed_img = \
        np.ascontiguousarray(np.transpose(preprocessed_img, (2, 0, 1)))
    preprocessed_img = torch.from_numpy(preprocessed_img)
    preprocessed_img.unsqueeze_(0)
    input = preprocessed_img.requires_grad_(True)
    return input


def show_cam_on_image(img, mask):
    heatmap = cv2.applyColorMap(np.uint8(255 * mask), cv2.COLORMAP_JET)
    heatmap = np.float32(heatmap) / 255
    cam = heatmap + np.float32(img)
    cam = cam / np.max(cam)
    cv2.imwrite("cam.jpg", np.uint8(255 * cam))


class GradCam:
    def __init__(self, model, feature_module, target_layer_names, use_cuda):
        self.model = model
        self.feature_module = feature_module
        self.model.eval()
        self.cuda = use_cuda
        if self.cuda:
            self.model = model.cuda()

        self.extractor = ModelOutputs(self.model, self.feature_module, target_layer_names)

    def forward(self, input):
        return self.model(input)

    def __call__(self, input, index=None):
        if self.cuda:
            features, output = self.extractor(input.cuda())  # [1, 2048, 7, 7], [1, 1000]
        else:
            features, output = self.extractor(input)

        # if index == None:
        #     index = np.argmax(output.cpu().data.numpy())
        # print(index)
        # print(output.size())
        # one_hot = np.zeros((1, (256,256)), dtype=np.float32)  # 1,1000
        # one_hot[0][index] = 1
        # one_hot = torch.from_numpy(one_hot).requires_grad_(True)
        output=torch.where(output>0.5,output,torch.full_like(output, 0))

        if self.cuda:
            one_hot = torch.sum(output)
            #print(one_hot)
        else:
            one_hot = torch.sum(output)

        self.feature_module.zero_grad()
        self.model.zero_grad()
        one_hot.backward(retain_graph=True)

        grads_val = self.extractor.get_gradients()[-1].cpu().data.numpy()  # 1, 2048, 7, 7
        target = features[-1]  # 1,2048,7,7
        target = target.cpu().data.numpy()[0, :]  # 2048, 7, 7

        weights = np.mean(grads_val, axis=(2, 3))[0, :]  # 2048
        cam = np.zeros(target.shape[1:], dtype=np.float32)

        for i, w in enumerate(weights):  # w:weight,target:feature
            cam += w * target[i, :, :]

        cam = np.maximum(cam, 0)  # 7,7
        cam = cv2.resize(cam, input.shape[2:])  # 224,224
        cam = cam - np.min(cam)
        cam = cam / np.max(cam)
        return cam




if __name__ == '__main__':
    """ python grad_cam.py <path_to_image>
    1. Loads an image with opencv.
    2. Preprocesses it for VGG19 and converts to a pytorch variable.
    3. Makes a forward pass to find the category index with the highest score,
    and computes intermediate activations.
    Makes the visualization. """

    grad_cam = GradCam(model=net, feature_module=model.outc, \
                       target_layer_names=["0"], use_cuda=args.use_cuda)
    # grad_cam = GradCam(model=model, feature_module=model.vgg1, \
    #                    target_layer_names=["29"], use_cuda=args.use_cuda)

    img = cv2.imread(args.image_path, 1)
    img = np.float32(cv2.resize(img, (256, 256))) / 255
    input = preprocess_image(img)

    # If None, returns the map for the highest scoring category.
    # Otherwise, targets the requested index.
    target_index = None
    mask = grad_cam(input, target_index)
    #print(mask.shape)

    show_cam_on_image(img, mask)














